package chat

import (
	"context"
	"fmt"
	"gitee.com/Cauchy_AQ/vxopen/rpc/transform/transformer"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/zeromicro/go-zero/core/logx"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512

	// send buffer size
	bufSize = 256
)

var (
	newline = []byte{'\n'}
	space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源，您可能需要根据实际情况修改这个检查
	},
}

type Connection struct {
	conn *websocket.Conn
	mtx  sync.Mutex
}

type Hub struct {
	connections map[*Connection]bool
	broadcast   chan []byte
	register    chan *Connection
	unregister  chan *Connection
	Transformer transformer.Transformer
}

func NewHub(transformers transformer.Transformer) *Hub {
	return &Hub{
		connections: make(map[*Connection]bool),
		broadcast:   make(chan []byte),
		register:    make(chan *Connection),
		unregister:  make(chan *Connection),
		Transformer: transformers,
	}
}

func (h *Hub) Run() {
	for {
		select {
		case conn := <-h.register:
			h.connections[conn] = true
			go h.AfterConnHandler(conn)

		case conn := <-h.unregister:
			if _, ok := h.connections[conn]; ok {
				delete(h.connections, conn)
				conn.conn.Close()
			}
		case message := <-h.broadcast:
			for conn := range h.connections {
				conn.Write(message)
			}
		}
	}
}

func (c *Connection) Write(message []byte) {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	err := c.conn.WriteMessage(websocket.TextMessage, message)
	if err != nil {
		logx.Errorf("Error writing message: %v", err)
	}
}

func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logx.Errorf("Error upgrading to WebSocket: %v", err)
		return
	}

	c := &Connection{conn: conn}
	h.register <- c

	defer func() {
		h.unregister <- c
	}()

	for {
		_, message, err := conn.ReadMessage()

		rpcResp, err1 := h.Transformer.WsMessageHandler(context.Background(), &transformer.WsMessageReq{
			Message: message,
		})
		if err1 != nil {
			logx.Errorf("Error ws rpc handle message: %v", err)
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				logx.Errorf("Error reading message: %v", err)
			}
			break
		}
		h.broadcast <- rpcResp.Message
	}
}

// AfterConnHandler 新连接上行，回发历史数据
func (h *Hub) AfterConnHandler(conn *Connection) {
	rpcResp, err := h.Transformer.WsConnectionHandler(context.Background(), &transformer.WsConnReq{
		Timestamp: time.Now().Unix(),
	})
	if err != nil {
		logx.Errorf("Error AfterConnHandler: %v", err)
	}
	for _, v := range rpcResp.Message {
		conn.Write(v)
	}
}

func init() {
	fmt.Println("Starting ws server at ws uri: '/chat' ...")
}
